# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of Cybertron package.
#
# The Cybertron is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Radical basis functions (RBF)"""

from typing import Union, Dict, Type, Optional

import torch
from torch import nn, Tensor
import numpy as np

from ..utils import get_integer, get_tensor, GLOBAL_DEVICE
from ..units import Units, GLOBAL_UNITS, Length, get_length


# Registry for RBF classes
_RBF_BY_KEY = {}


def _rbf_register(*aliases):
    """Return the alias register."""
    def alias_reg(cls):
        name = cls.__name__.lower()
        if name not in _RBF_BY_KEY:
            _RBF_BY_KEY[name] = cls

        for alias in aliases:
            if alias not in _RBF_BY_KEY:
                _RBF_BY_KEY[alias] = cls
        return cls
    
    return alias_reg


# Define the RadicalBasisFunctions class
class RadicalBasisFunctions(nn.Module):
    """
    Base class for radical basis functions.
    
    Args:
        num_basis (int): Number of basis functions.
        r_max (float or torch.Tensor): Maximum distance for RBF.
        r_min (float or torch.Tensor, optional): Minimum distance for RBF. Defaults to 0.
        clip_distance (bool, optional): Whether to clip distance to [r_min, r_max]. Defaults to False.
        length_unit (str, optional): Unit of length. Defaults to 'nm'.
    """
    def __init__(self, 
                 num_basis: int, 
                 r_max: Union[Length, float, Tensor], 
                 r_min: Union[Length, float, Tensor] = 0, 
                 clip_distance: bool = False, 
                 length_unit: str = 'nm',
                 **kwargs
                 ) -> None:
        super().__init__()
        self._kwargs = kwargs

        self.num_basis = get_integer(num_basis)
        self.device = GLOBAL_DEVICE()

        if length_unit is None:
            length_unit = GLOBAL_UNITS.length_unit
        self.units = Units(length_unit)

        self.r_max = get_tensor(r_max, dtype=torch.float32, device=self.device)
        self.r_min = get_tensor(r_min, dtype=torch.float32, device=self.device)
        self.clip_distance = clip_distance
        self.length_unit = length_unit

        if self.length_unit is None:
            self.length_unit = GLOBAL_UNITS.length_unit

        if self.r_max <= self.r_min:
            raise ValueError(f'In RBF, r_max ({self.r_max}) must be larger than r_min ({self.r_min})!')

        self.r_range = self.r_max - self.r_min

    def _print_info(self,
                    num_retraction: int = 6,
                    num_gap: int = 3,
                    char: str = ' '
                    ) -> None:
        ret = char * num_retraction
        gap = char * num_gap
        print(ret+gap+f' Minimum distance: {self.r_min} {self.length_unit}')
        print(ret+gap+f' Maximum distance: {self.r_max} {self.length_unit}')
        print(ret+gap+f' Number of basis functions: {self.num_basis}')
        if self.clip_distance:
            print(ret+gap+f' Clip the range of distance to ({self.r_min}, {self.r_max}).')

    def forward(self, distance: Tensor) -> Tensor:
        """Forward pass to compute RBF values."""
        return self.compute_rbf(distance)

    def compute_rbf(self, distance: Tensor) -> Tensor:
        """Compute RBF values for given distances."""
        raise NotImplementedError

    def __str__(self):
        return 'RadicalBasisFunctions<>'